#!/usr/bin/python3
# -*- coding: utf-8 -*-

import numpy  as np
import pandas as pd
import talib as ta

def signal(*args):
    # MtmMean 指标 排序顺序改为False，False是指从大到小排序，也就是说涨得最快的最多，跌得最猛的做空
    df = args[0]
    n = args[1]
    # diff_num = args[2]
    # factor_name = args[3]
    factor_name = args[2]

    # df[factor_name] = (df['close'] / df['close'].shift(n) - 1).rolling(
    #     window=n, min_periods=1).mean()
    pct_change = df['close'].pct_change(n)
    df['MtmMean'] = pct_change.rolling(n).mean()
    df['std'] = pct_change.rolling(n).std()
    df[factor_name] = (df['MtmMean'] * df['std'])

    '''
    # 计算价格变化的百分比一次，重复使用结果
pct_change = df['close'].pct_change(n)

# 使用assign方法一次性计算多个新列
df = df.assign(
    MtmMean=pct_change.rolling(n).mean(),
    std=pct_change.rolling(n).std()
).assign(
    **{factor_name: lambda x: x['MtmMean'] * x['std']}
)
    '''

    return df

